import os
import numpy as np
import cv2
import colorsys
from time import time
import rclpy
from rclpy.node import Node
from sensor_msgs.msg import CompressedImage
import cv2
import numpy as np  
from std_msgs.msg import Int32MultiArray
from std_msgs.msg import Int32
from ai_msgs.msg import PerceptionTargets
from origincar_msg.msg import Sign  # 根据图片的消息类型
from origincar_msg.msg import Data
from std_msgs.msg import String
from openai import OpenAI
import base64
import json
import subprocess
from concurrent.futures import ThreadPoolExecutor, wait, FIRST_COMPLETED
import time
from pyzbar.pyzbar import decode 
import argparse

class FcosDetection(Node):
    def __init__(self,qr_choice=0):
        super().__init__('fcos_detection')

        # 初始化微信二维码检测器
        self.wechat_qr_detector = WeChatQRCodeDetector()

        self.qr_choice = qr_choice  # 0为pyzbar解码，1为wechat_qr解码

        self.qrcode_info = 0
        self.task = 0

        #0为默认，1为启动大模型调用
        self.vlm_flag = 0
        #0为默认，1为启动cv检测
        self.opencv_flag = 0

        self.park_relativ_position = np.array([0.0, 0.0,0.0,0.0])
        self.zhuitong_relativ_position = np.array([0.0, 0.0,0.0,0.0])
        self.task1_time = 0
        self.task3_time = 0

        # 创建一个订阅者，订阅 /sigh4return 话题，消息类型为 Int32,上位机信号
        self.sign4return = self.create_subscription(
            Int32,
            '/sign4return',
            self.sign4return_callback,
            10)     
        # 创建订阅者，订阅话题 /image
        self.subscription = self.create_subscription(
            CompressedImage,
            '/image',
            self.image_callback,
            10)
        self.subscription_vlm = self.create_subscription(  
            Int32,  
            '/vlm',  
            self.vlm_callback,  
            10)
            
        #创建订阅者，订阅话题 hobot_dnn_detection 话题   
        self.subscription_fcos = self.create_subscription( 
            PerceptionTargets, 
            '/hobot_dnn_detection',  
            self.fcos_callback, 
            10)  


        # 创建发布者，发布到 /sign_switch 话题
        self.sign_switch_pub = self.create_publisher(Sign, 'sign_switch', 10)      # 发布二维码的信息
        #创建新的发布者，用于发布避障信息
        self.cone_publisher = self.create_publisher(Int32MultiArray, '/cone', 10)  # 创建新的发布者，用于发布x和y坐标
        #创建新的发布者，用于发布park信息
        self.park_publisher = self.create_publisher(Int32MultiArray, '/park', 10)  # 创建新的发布者，用于发布x和y坐标   
        #创建发布者，发布vlm识别结果
        self.vlm_msg_publisher = self.create_publisher(String, '/vlm_msg', 10)
        #创建发布者，发布二维码识别结果
        self.qrcode_publisher = self.create_publisher(String, '/qrcode', 10)
        #创建发布者，发布fcos识别结果图像
        #self.image_fcos_publisher = self.create_publisher(CompressedImage,'/image_fcos',10)

        #self.bridge = CvBridge()

    def image_callback(self, msg):
        
        #将图像数据转换为OpenCV格式
        np_arr = np.frombuffer(msg.data, np.uint8)
        frame = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)

        if(self.vlm_flag == 1): #云端大模型识别
            self.get_logger().info(f"开始启动大模型")
            start_time = self.get_clock().now()
            vlm_msg = String()
            #vlm_msg.data = self.qwen_vl(frame=cv_image)
            vlm_msg.data = str(self.qwen_vl(frame=frame))
            self.vlm_msg_publisher.publish(vlm_msg)
            self.get_logger().info(f"大模型识别结果为: {vlm_msg.data}")
            elapsed = self.get_clock().now() - start_time
            elapsed_sec = elapsed.nanoseconds / 1e9
            self.get_logger().info(f"✅ 使用的是【云端大模型】，识别耗时：{elapsed_sec:.2f} 秒")
            self.vlm_flag = 0
        
        if(self.vlm_flag == 2): #本地/云端大模型识别
            image_path = "/root/vlm.jpg"
            cv2.imwrite(image_path, frame)
            self.get_logger().info(f"开始启动大模型")
            executor = ThreadPoolExecutor(max_workers=2)
            start_time = self.get_clock().now()

            future_cloud = executor.submit(self.safe_qwen_vl, frame)
            future_local = executor.submit(self.run_local_vlm, image_path, "描述图片里的图片/照片（是个病人，请只描述病人）")

            done, _ = wait([future_cloud, future_local], return_when=FIRST_COMPLETED)
            first = done.pop()

            result = first.result()
            if result is None:
                self.get_logger().warn("⚠️ 云端模型不可用，等待本地模型结果...")
                second = (set([future_cloud, future_local]) - set([first])).pop()
                result = second.result()
                source = "本地大模型" if second == future_local else "云端大模型"
            else:
                source = "云端大模型" if first == future_cloud else "本地大模型"

            elapsed = self.get_clock().now() - start_time
            elapsed_sec = elapsed.nanoseconds / 1e9

            # 尝试主动关闭线程池，不等待慢线程（让其后台结束）
            executor.shutdown(wait=False)

            # 发布结果
            vlm_msg = String()
            vlm_msg.data = str(result)
            self.vlm_msg_publisher.publish(vlm_msg)

            self.get_logger().info(f"✅ 使用的是【{source}】，识别耗时：{elapsed_sec:.2f} 秒")
            self.get_logger().info(f"大模型识别结果为: {vlm_msg.data}")

            self.vlm_flag = 0

        if(self.task == 1): #二维码识别
            time = self.get_clock().now()
            dt_1 = (time - self.task1_time).nanoseconds * 1e-9
            #if dt_1 > 1.5:
            if self.qrcode_info == 0:
                if self.qr_choice == 0:#确保只打印一次
                    self.get_logger().info(f"使用pyzbar解码") 
                else:
                    self.get_logger().info(f"使用wechat_qr解码")
                self.qrcode_info = 1
                #image_roi = cv_image[0:480,0:640] #截取图片，先高后宽 [50:450,0:640]
            if self.qr_choice == 0:  # 使用 pyzbar 解码二维码
                self.decode_qr_code_pyzbar(frame)  # 使用 pyzbar 解码二维码
            else:
                self.decode_qr_code_wechat(frame)  # 使用 WeChatQRCode 解码二维码   

    def fcos_callback(self, msg):
        for num, target in enumerate(msg.targets):
            
            if target.rois:
                
                class_name = target.rois[0].type.strip()
                # 获取第一个 ROI
                roi = target.rois[0].rect

                # 获取框参数
                x_offset = roi.x_offset
                y_offset = roi.y_offset
                height = roi.height
                width = roi.width
                confidence =target.rois[0].confidence
                
            if (self.task == 1) :
                if class_name == 'obstacle': 
                    #底边中点坐标
                    x = int(x_offset + 0.5 * width)
                    y = y_offset + height
                    self.zhuitong_relativ_position = x,y,width,height
                    if (width > 90 and height > 100) : #只发布近处需要避障的锥桶，屏蔽远处锥桶带来的干扰
                        #把坐标发布出去方便调试
                        cone_msg = Int32MultiArray()
                        cone_msg.data = self.zhuitong_relativ_position
                        self.cone_publisher.publish(cone_msg)
                else:
                     #把坐标发布出去方便调试
                     cone_msg = Int32MultiArray()
                     cone_msg.data = 0,0,0,0
                     self.cone_publisher.publish(cone_msg)

                #if class_name == 'qr':


                
            
            if(self.opencv_flag == 1):    #按键k开启fcos检测，便于观察识别结果

                if class_name == 'obstacle': 
                    #底边中点坐标
                    x = int(x_offset + 0.5 * width)
                    y = y_offset + height
                    self.zhuitong_relativ_position = x,y,width,height
                    cone_msg = Int32MultiArray()
                    cone_msg.data = self.zhuitong_relativ_position
                    self.cone_publisher.publish(cone_msg)

                if class_name == 'p':
                    x = int(x_offset + 0.5 * width)
                    y = int(y_offset + 0.5 * height)
                    self.park_relativ_position = x,y,width,height
                    park_msg = Int32MultiArray()
                    park_msg.data = self.park_relativ_position
                    self.park_publisher.publish(park_msg)

                if class_name != 'obstacle': 
                    cone_msg = Int32MultiArray()
                    cone_msg.data = 0,0,0,0
                    self.cone_publisher.publish(cone_msg)

                if class_name != 'p':
                    park_msg = Int32MultiArray()
                    park_msg.data = 0,0,0,0
                    self.park_publisher.publish(park_msg)

    def vlm_callback(self, msg):  #接受按键信息，并设置vlm_flag为1，启动大模型，可设置cv检测
        self.vlm = msg.data
        if self.vlm == 1:
            self.vlm_flag = 1
        if self.vlm == 2:
            self.opencv_flag = 1
        if self.vlm == 3:
            self.opencv_flag = 0
        if self.vlm == 4:
            self.vlm_flag = 2

    def qwen_vl(self,frame):    #调用qwen API

        client = OpenAI(
            # 若没有配置环境变量，请用百炼API Key将下行替换为：api_key="sk-xxx",
            api_key=os.getenv("DASHSCOPE_API_KEY"),
            base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
        )
        #  Base64 编码格式
        # def encode_image(image_path):
        #     with open(image_path, "rb") as image_file:
        #         return base64.b64encode(image_file.read()).decode("utf-8")

        """将 OpenCV 读取的帧转换为 Base64"""
        _, buffer = cv2.imencode(".png", frame)  # 将 frame 转换为 PNG 格式的字节流
        base64_image = base64.b64encode(buffer).decode("utf-8")
            
        completion = client.chat.completions.create(
            model="qwen-vl-plus",  # 此处以qwen-vl-plus为例，可按需更换模型名称。模型列表：https://help.aliyun.com/zh/model-studio/getting-started/models
            messages=[{"role": "user","content": [
                    {"type": "text","text": "中文描述图片内的插图或者说海报或者说白色立牌上病人的内容"},
                    {"type": "image_url",
                    #"image_url": {"url": "https://dashscope.oss-cn-beijing.aliyuncs.com/images/dog_and_girl.jpeg"}}
                    "image_url": {"url": f"data:image/png;base64,{base64_image}"}}
                    ]}],
            #max_tokens = 30,
            )
        response_str = completion.model_dump_json()
        response_dict = json.loads(response_str)
        content = response_dict["choices"][0]["message"]["content"]

        return content

    def safe_qwen_vl(self,frame):
        try:
            return self.qwen_vl(frame)
        except Exception as e:
            self.get_logger().warn(f"⚠️ 云端模型调用失败：{e}")
            return None

    def sign4return_callback(self,msg): #sign4return话题的回调函数，在foxglove上显示，传递self.task的值

    # 创建 Sign 消息实例
        if msg.data == 6:

            self.task = 3
            self.task3_time = self.get_clock().now()

        if msg.data == -1:

            self.task = 1
            self.task1_time = self.get_clock().now()
            
        if msg.data == 5:
            self.task = 2

        return
    
    def publish_sign_message(self, qr_code_data,sign_data): #发布二维码的信息
        '''
        qr_code_data: 识别到的二维码数据
        sign_data: 识别到的二维码数据对应的sign数据 3或4
        '''
        # 创建并发布 Sign 消息
        sign_msg = Sign()
        sign_msg.sign_data = sign_data
        self.sign_switch_pub.publish(sign_msg)
        qrcode_msg = String()
        qrcode_msg.data =  f"二维码的识别结果为: {qr_code_data}"
        self.qrcode_publisher.publish(qrcode_msg)
        self.get_logger().info(f"已发布二维码信息: {qr_code_data}")

    def decode_qr_code_pyzbar(self, image):  #识别到二维码结果后停车，进入任务二
        # 使用 pyzbar 进行二维码解码
        #image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        #image = cv2.resize(image,(320,240)) #处理更快，但是精度更低
        decoded_objects = decode(image)
        
        #decoded_objects = decode(image, symbols=[ZBarSymbol.QRCODE])
        if decoded_objects:
            for obj in decoded_objects:
                qr_code_data = obj.data.decode('utf-8') #从字节码转换为字符串
                if qr_code_data in ["ClockWise", "AntiClockWise"]:
                    sign_data = 3 if qr_code_data == "ClockWise" else 4
                    self.publish_sign_message(qr_code_data,sign_data)
                    self.task = 2
                    return
                else:
                    qr_code_data = int(qr_code_data)
                    sign_data = 3 if qr_code_data % 2 == 1 else 4
                    self.publish_sign_message(qr_code_data,sign_data)
                    self.task = 2
                    return  # 成功识别并发布后立即返回，避免重复处理
            else:
                self.get_logger().info(f"Non-numeric QR code detected")
        else:
            self.get_logger().info(f"No QR code detected.")

    def decode_qr_code_wechat(self, image):
        # 使用微信二维码检测器进行解码
        qr_code_data = self.wechat_qr_detector.detect(image)

        if qr_code_data:
            if qr_code_data in ["ClockWise", "AntiClockWise"]:
                sign_data = 3 if qr_code_data == "ClockWise" else 4
                self.publish_sign_message(qr_code_data,sign_data)
                self.task = 2
                return
            else:
                qr_code_data = int(qr_code_data)
                sign_data = 3 if qr_code_data % 2 == 1 else 4
                self.publish_sign_message(qr_code_data,sign_data)
                self.task = 2
                return  # 成功识别并发布后立即返回，避免重复处理
        else:
            self.get_logger().info(f"No QR code detected.")

    def run_local_vlm(self,image_path, prompt):
        cmd = [
            "/root/llama.cpp_internvl2_bpu/llama.cpp/build/bin/llama-intern2vl-bpu-cli",
            "-m", "/root/llama.cpp_internvl2_bpu/llama.cpp/Qwen2.5-0.5B-Instruct-Q4_0.gguf",
            "--mmproj", "/root/llama.cpp_internvl2_bpu/llama.cpp/vit_model_int16_v2.bin",
            "--image", image_path,
            "-p", prompt,
            "--temp", "0.5",
            "--threads", "8"
        ]
        try:
            result = subprocess.run(cmd, capture_output=True, text=True, check=True)
            # 解析输出，只保留最后一句模型输出
            lines = result.stdout.strip().splitlines()
            output = lines[-1] if lines else "（无输出）"
            return output
        except subprocess.CalledProcessError as e:
            return f"模型运行失败: {e.stderr}"

class WeChatQRCodeDetector:
    def __init__(self):
        self.depro = '/root/lyz_qrcode/detect.prototxt'
        self.decaf = '/root/lyz_qrcode/detect.caffemodel'
        self.srpro = '/root/lyz_qrcode/sr.prototxt'
        self.srcaf = '/root/lyz_qrcode/sr.caffemodel'

    def detect(self, img):
        #gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        detector = cv2.wechat_qrcode_WeChatQRCode(self.depro, self.decaf, self.srpro, self.srcaf)
        barcodes, _ = detector.detectAndDecode(img)
        return barcodes[0] if barcodes else None

def main(args=None):  
    rclpy.init(args=args)  # 初始化rclpy  
    parser = argparse.ArgumentParser(description="FCOS Detection with optional QR code decoder.")
    parser.add_argument('--qr_choice', type=int, default=0, choices=[0, 1],
                        help='QR code decoder choice: 0=pyzbar, 1=wechat_qr')
    args = parser.parse_args()
    # 创建节点对象时传递参数
    fcosdetection = FcosDetection(qr_choice=args.qr_choice)
    rclpy.spin(fcosdetection)
    fcosdetection.destroy_node()  # 销毁节点
    rclpy.shutdown()
  
if __name__ == '__main__':  
    main()  # 入口函数
